"""
This is part of Human Robot Dynamics-Control, By Tourists.

Based on "OpenLoong Dynamics Control"
Transfer from mujoco to ISAAC SIM by IsaacLab

"""

import numpy as np
from scipy import linalg

## SVD based pseudo-inverse
def pseudoInv_SVD(MAT = np.array([])):
    # auto svd = mat.jacobiSvd(Eigen::ComputeFullU | Eigen::ComputeFullV);
    # 计算奇异值分解
    u, s, vh = np.linalg.svd(MAT, full_matrices=True)
    maxSgVal = 0.0 
    singularValuesInv = np.array([])
    singular_values = np.diag(s)
    max_singular_value = np.max(np.abs(s))
    
    tolerance = np.finfo(float).eps * max(MAT.shape) * max_singular_value
    singular_values_inv = np.zeros_like(singular_values)
    for i in range(singular_values.shape[0]):
        if np.abs(singular_values[i, i]) > tolerance:
            singular_values_inv[i, i] = 1 / singular_values[i, i]
    return vh.T @ singular_values_inv @ u.T
    

def pseudoInv_right(M = np.array([])):
    Mres = np.array([])
    Mres = M @ M.T
    I = np.eye(Mres.shape[0])
    # Mres=M.transpose()* Mres.ldlt().solve(Eigen::MatrixXd::Identity(Mres.rows(), Mres.cols()));
    Mres = M.T @ np.linalg.solve(Mres, I)
    ## Mres=M.transpose()* Mres.inverse();
    return Mres

# def pseudoInv_right_weighted(M, W):
#     damp = 0
#     # Mres = M @ np.diag(1 / W) @ M.T
#     Mres = M @ np.linalg.inv(W) @ M.T
#     Mres[np.diag_indices_from(Mres)] += damp
#     U, S, Vt = np.linalg.svd(Mres)
#     S_inv = np.zeros(Mres.shape).T
#     S_inv[:len(S), :len(S)] = np.diag(1 / S)
#     Mres = np.diag(1 / W) @ M.T @ (U @ S_inv @ Vt)
#     return Mres

def pseudoInv_right_weighted(M, W, eps=1e-10):
    """
    Compute the right-weighted pseudo-inverse of M.
    
    Args:
    M: Input matrix
    W: Weight matrix (should be a 1D array representing the diagonal)
    eps: Small value for numerical stability
    
    Returns:
    Right-weighted pseudo-inverse of M
    """
    damp = 0
    # Ensure W is a 1D array
    if W.ndim > 1:
        W = np.diag(W) 
    # Compute W_inv carefully
    W_inv = np.diag(1 / (W + eps))
    # Compute Mres
    Mres = M @ W_inv @ M.T
    # Add damp to the diagonal
    np.fill_diagonal(Mres, Mres.diagonal() + damp)
    # Use QR decomposition with pivoting for pseudo-inverse (similar to completeOrthogonalDecomposition in Eigen)
    Q, R, P = linalg.qr(Mres, pivoting=True, mode='economic')
    # Compute rank and truncate
    tol = eps * np.abs(R[0, 0])
    rank = np.sum(np.abs(np.diag(R)) > tol)
    Q = Q[:, :rank]
    R = R[:rank, :rank]
    # Compute pseudo-inverse
    Rinv = linalg.inv(R)
    Mres_pinv = Q @ Rinv @ Rinv.T @ Q.T
    # Compute final result
    result = W_inv @ M.T @ Mres_pinv
    return result

# def dyn_pseudo_inv(M, dyn_A):
#     #--------------------------
#     res = None
#     Minv = inv(dyn_A)
#     res = Minv @ M.T @ inv(M @ Minv @ M.T)
#     return res

def dyn_pseudoInv(M, dyn_M, isMinv):
    damp = 0
    Minv = None

    if isMinv:
        Minv = dyn_M
    else:
        Minv = np.linalg.inv(dyn_M)

    temp = M @ Minv @ M.T

    np.fill_diagonal(temp, np.diagonal(temp) + damp)

    res = Minv @ M.T @ np.linalg.pinv(temp)

    return res

def eul2Rot(roll, pitch, yaw):
    Rz = np.array([[np.cos(yaw), -np.sin(yaw), 0],
                   [np.sin(yaw), np.cos(yaw), 0],
                   [0, 0, 1]])
    Ry = np.array([[np.cos(pitch), 0, np.sin(pitch)],
                   [0, 1, 0],
                   [-np.sin(pitch), 0, np.cos(pitch)]])
    Rx = np.array([[1, 0, 0],
                   [0, np.cos(roll), -np.sin(roll)],
                   [0, np.sin(roll), np.cos(roll)]])
    return Rz @ Ry @ Rx

def rot2eul(rot):
    eul = np.zeros(3)
    eul[0] = np.arctan2(rot[2, 1], rot[2, 2])
    eul[1] = np.arctan2(-rot[2, 0], np.sqrt(rot[2, 1] * rot[2, 1] + rot[2, 2] * rot[2, 2]))
    eul[2] = np.arctan2(rot[1, 0], rot[0, 0])
    return eul

def eigen2std(eigenVector):
    return list(eigenVector)

def eul2quat(roll, pitch, yaw):
    R = eul2Rot(roll, pitch, yaw)
    # quatCur = R.from_euler()
    quatCur = R; ## rotation matrix converted to quaternion
    # Eigen::Quaterniond resQuat;
    r = R.from_euler('xyz', [roll, pitch, yaw], degrees=True)
    return r.as_quat()


# def diffRot(Rcur, Rdes):
#     R = np.dot(Rcur.T, Rdes)
#     w = np.zeros(3)

#     if np.allclose(np.diag(R), np.ones(3), atol=1e-5) and np.abs(np.trace(R) - 3) < 1e-3:
#         w.fill(0)
#     elif np.allclose(np.diag(R), np.ones(3), atol=1e-5):
#         # w = (R[0, 0] + 1, R[1, 1] + 1, R[2, 2] + 1)
#         # w = np.array(w) * np.pi / 2.0
#         w = np.array([R[0, 0] + 1, R[1, 1] + 1, R[2, 2] + 1]) * np.pi / 2.0
#     else:
#         l = np.array([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]])
#         sita = np.arctan2(np.linalg.norm(l), np.trace(R) - 1)
#         # w = sita * l / np.linalg.norm(l)
#         if np.linalg.norm(l) != 0:
#             w = sita * l / np.linalg.norm(l)
#         else:
#             w = np.zeros(3)

#     w = np.dot(Rcur, w)
#     return w

def diffRot(Rcur, Rdes):
    R = Rcur.T @ Rdes
    
    # check ones
    if np.allclose(np.diag(R), [1, 1, 1], atol=1e-5) and np.sum(np.abs(np.diag(R))) - 3 < 1e-3:
        w = np.zeros(3)
    elif np.allclose(np.diag(R), [1, 1, 1], atol=1e-5):
        w = np.array([R[0, 0] + 1, R[1, 1] + 1, R[2, 2] + 1]) * (3.1415 / 2.0)
    else:
        l = np.array([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]])
        sita = np.arctan2(np.linalg.norm(l), R[0, 0] + R[1, 1] + R[2, 2] - 1)
        w = sita * l / np.linalg.norm(l)
    
    w = Rcur @ w
    return w

def quat2axisAngle(quat):
    # quat must be normalized
    res = np.zeros(4)
    angle = 2.0 * np.arccos(quat[3])
    s = np.sqrt(1 - quat[3] * quat[3])
    if s < 1e-8:
        res[0] = quat[0]
        res[1] = quat[1]
        res[2] = 1  # no rotation now but the axis must not be [0,0,0]
    else:
        res[0] = quat[0] / s
        res[1] = quat[1] / s
        res[2] = quat[2] / s
    res[3] = angle
    return res

from scipy.spatial.transform import Rotation as R
def intQuat(quat, w):
    r = R.from_quat(quat)
    Rcur = r.as_matrix()
    w_norm = w / np.linalg.norm(w)
    if np.linalg.norm(w) > 1e-4:
        a = np.array([[0, -w_norm[2], w_norm[1]],
                      [w_norm[0], 0, -w_norm[0]],
                      [-w_norm[1], w_norm[0], 0]])
        theta = np.linalg.norm(w)
        Rinc = np.identity(3) + a * np.sin(theta) + np.dot(a, a) * (1 - np.cos(theta))
    else:
        Rinc = np.identity(3)
    Rend = np.dot(Rcur, Rinc)
    quatRes = R.from_matrix(Rend).as_quat()
    return quatRes

def sign(x):
    if x > 0:
        return 1.0
    elif x < 0:
        return -1.0
    else:
        return 0


def Rx3(theta):  # local to world
    # for 2D-XY vector, rotation matrix along z axis
    M = np.zeros((3, 3))
    M[0, 0] = 1
    M[1, 1] = np.cos(theta)
    M[1, 2] = -np.sin(theta)
    M[2, 1] = np.sin(theta)
    M[2, 2] = np.cos(theta)
    return M

def Ry3(theta):  # local to world
    # for 2D-XY vector, rotation matrix along z axis
    M = np.zeros((3, 3))
    M[0, 0] = np.cos(theta)
    M[0, 1] = 0
    M[0, 2] = np.sin(theta)
    M[1, 1] = 1
    M[1, 2] = 0
    M[2, 0] = -np.sin(theta)
    M[2, 1] = 0
    M[2, 2] = np.cos(theta)
    return M

def Rz3(theta):  # local to world
    # for 2D-XY vector, rotation matrix along z axis
    M = np.zeros((3, 3))
    M[0, 0] = np.cos(theta)
    M[0, 1] = -np.sin(theta)
    M[0, 2] = 0
    M[1, 0] = np.sin(theta)
    M[1, 1] = np.cos(theta)
    M[1, 2] = 0
    M[2, 0] = 0
    M[2, 1] = 0
    M[2, 2] = 1
    return M

def cross_product_A(A):
    M = np.zeros((3, 3))
    M[0, 1] = -A[2]
    M[0, 2] = A[1]
    M[1, 0] = A[2]
    M[1, 2] = -A[0]
    M[2, 0] = -A[1]
    M[2, 1] = A[0]
    return M

def Ramp(u, tgt, inc):
    output = 0.0
    if abs(u - tgt) < inc:
        output = tgt
    elif u < tgt - inc:
        output = u + inc
    elif u > tgt + inc:
        output = u - inc
    else:
        output = tgt

    return output

def limit(data, max_value, min_value):
    if data > max_value:
        data = max_value
    if data < min_value:
        data = min_value
    return data
